{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert LAMMPS DATA file to VASP POSCAR file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "print('This script converts a LAMMPS DATA file into a VASP POSCAR file.')\n",
    "print('This script is interactive and requires user input. Please read carefully before proceeding:\\n')\n",
    "\n",
    "print('The lattice may need to be reoptimized using DFT. As such, all atoms are printed in direct coordinates.')\n",
    "print('For reverse conversion, use poscar2lmp.awk\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract unit cell information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwd = os.getcwd()\n",
    "\n",
    "name = input('Enter the name of the LAMMPS DATA file to be converted : ')\n",
    "file = pwd + '/' + name\n",
    "log = open(file,'r')\n",
    "log_lines = log.readlines()\n",
    "log_lines = [line.split() for line in log_lines]\n",
    "\n",
    "for line_index, line in enumerate(log_lines):\n",
    "    if line:\n",
    "        for l in line:\n",
    "            \n",
    "            if l == 'atoms':\n",
    "                N_atom = int(line[0])    # Number of atoms\n",
    "            \n",
    "            elif l == 'xlo':\n",
    "                xlo = float(line[0])    # Box dimensions\n",
    "                xhi = float(line[1])\n",
    "            \n",
    "            elif l == 'ylo':\n",
    "                ylo = float(line[0])\n",
    "                yhi = float(line[1])\n",
    "            \n",
    "            elif l == 'zlo':\n",
    "                zlo = float(line[0])\n",
    "                zhi = float(line[1])\n",
    "            \n",
    "            elif l == 'xy':\n",
    "                xy = float(line[0])\n",
    "                xz = float(line[1])\n",
    "                yz = float(line[2])\n",
    "            \n",
    "            elif l == 'Atoms':\n",
    "                start = line_index + 2\n",
    "            \n",
    "            elif l == 'Velocities':\n",
    "                end = line_index - 1\n",
    "\n",
    "for line in log_lines[start:end]:\n",
    "    line[:2] = np.array(line[:2]).astype(int)    # Atom ID, atom type\n",
    "    line[2:5] = np.array(line[2:5]).astype(float)    # x, y, z\n",
    "\n",
    "log.close()\n",
    "\n",
    "lat = np.array([[xhi-xlo, 0.0, 0.0], [xy, yhi-ylo, 0.0], [xz, yz, zhi-zlo]])    # Lattice matrix\n",
    "\n",
    "print('Unit cell information extracted.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract system information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys = input('Enter the system name : ')\n",
    "\n",
    "a = ' '.join(str(l) for l in lat[0])\n",
    "b = ' '.join(str(l) for l in lat[1])\n",
    "c = ' '.join(str(l) for l in lat[2])\n",
    "\n",
    "scale = 1.00000000000000    # Assumes no lattice scaling\n",
    "\n",
    "ls_element = input('Enter the list of the elements, separated by space, in the same order as LAMMPS atom types : ')\n",
    "\n",
    "ls_type = []    # List of atom types\n",
    "for line in log_lines[start:end]:\n",
    "    Type = line[1]\n",
    "    if Type not in ls_type:\n",
    "        ls_type.append(Type)\n",
    "N_element = len(ls_type)    # Number of elements\n",
    "\n",
    "N_type = []    # List of number of atoms of each atom type\n",
    "for atom_type in ls_type:\n",
    "    N = 0\n",
    "    for line in log_lines[start:end]:\n",
    "        Type = line[1]\n",
    "        if Type == atom_type:\n",
    "            N += 1\n",
    "    N_type.append(N)\n",
    "N_type = ' '.join(str(n) for n in N_type)\n",
    "\n",
    "F = 'F F F'    # Selective dynamics\n",
    "T = 'T T T'\n",
    "\n",
    "fix = input('Selective dynamics? Yes or No : ')\n",
    "if fix == 'Yes':\n",
    "    print('The fixed atoms are assumed to be numbered consecutively (e.g. height-by-height). Please quit now if not so.')\n",
    "    header = '%s\\n%f\\n%s\\n%s\\n%s\\n%s\\n%s\\n%s\\n%s\\n'%(sys, scale, a, b, c, ls_element, N_type, 'Selective dynamics', 'Direct')\n",
    "    fix_range = input('Enter the IDs of the first and last atom to be fixed, separated by space : ')\n",
    "    fix_range = fix_range.split()\n",
    "    fix_range = np.array(fix_range).astype(int)\n",
    "else:\n",
    "    header = '%s\\n%f\\n%s\\n%s\\n%s\\n%s\\n%s\\n%s\\n'%(sys, scale, a, b, c, ls_element, N_type, 'Direct')\n",
    "\n",
    "print('System information extracted.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to direct coordinates & print"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_poscar = pwd + '/POSCAR'\n",
    "poscar = open(file_poscar,'w')\n",
    "poscar.write(header)\n",
    "\n",
    "for atom_type in ls_type:    # List atoms element-by-element\n",
    "    for line in log_lines[start:end]:\n",
    "        \n",
    "        ID = line[0]\n",
    "        Type = line[1]\n",
    "        \n",
    "        if Type == atom_type:\n",
    "\n",
    "            x = line[2]\n",
    "            y = line[3]\n",
    "            z = line[4]\n",
    "            r = np.array([x, y, z])    # Cartesian coordinates\n",
    "            Rep = np.dot(r, np.linalg.inv(lat))    # Direct coordinates\n",
    "                        \n",
    "            if fix == 'Yes':    # Selective dynamics\n",
    "                if fix_range[0] <= ID <= fix_range[1]:\n",
    "                    coord = '%f %f %f %s\\n'%(Rep[0], Rep[1], Rep[2], F)\n",
    "                    poscar.write(coord)\n",
    "                else:\n",
    "                    coord = '%f %f %f %s\\n'%(Rep[0], Rep[1], Rep[2], T)\n",
    "                    poscar.write(coord)\n",
    "                            \n",
    "            else:\n",
    "                coord = '%f %f %f\\n'%(Rep[0], Rep[1], Rep[2])\n",
    "                poscar.write(coord)\n",
    "\n",
    "poscar.close()\n",
    "\n",
    "print('POSCAR generated > ./POSCAR')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
